"""
Mesh input/output operations
"""
import os
import numpy as np
import struct
try:
    import vtk
    from vtk.util.numpy_support import vtk_to_numpy, numpy_to_vtk
    VTK_AVAILABLE = True
except ImportError:
    VTK_AVAILABLE = False

class MeshIO:
    """Class for handling mesh input and output operations"""
    
    @staticmethod
    def read_mesh(file_path):
        """
        Read mesh from file
        
        Parameters:
        -----------
        file_path : str
            Path to the mesh file
            
        Returns:
        --------
        dict
            Dictionary containing mesh data (nodes, elements, etc.)
        """
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"Mesh file not found: {file_path}")
        
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == '.vtk':
            if VTK_AVAILABLE:
                return MeshIO._read_vtk(file_path)
            else:
                raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        elif ext == '.stl':
            if VTK_AVAILABLE:
                return MeshIO._read_stl(file_path)
            else:
                return MeshIO._read_stl_binary(file_path)
        else:
            raise ValueError(f"Unsupported file format: {ext}")
    
    @staticmethod
    def write_mesh(mesh_data, file_path):
        """
        Write mesh to file
        
        Parameters:
        -----------
        mesh_data : dict
            Dictionary containing mesh data
        file_path : str
            Path to save the mesh file
        """
        ext = os.path.splitext(file_path)[1].lower()
        
        if ext == '.vtk':
            if VTK_AVAILABLE:
                MeshIO._write_vtk(mesh_data, file_path)
            else:
                raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        elif ext == '.stl':
            if VTK_AVAILABLE and 'vtk_mesh' in mesh_data:
                MeshIO._write_stl(mesh_data, file_path)
            else:
                MeshIO._write_stl_binary(mesh_data, file_path)
        else:
            raise ValueError(f"Unsupported file format: {ext}")
    
    @staticmethod
    def _read_vtk(file_path):
        """Read VTK file"""
        if not VTK_AVAILABLE:
            raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        
        reader = vtk.vtkUnstructuredGridReader()
        reader.SetFileName(file_path)
        reader.Update()
        
        mesh = reader.GetOutput()
        
        # Extract points
        points = vtk_to_numpy(mesh.GetPoints().GetData())
        
        # Extract cells
        cells = []
        for i in range(mesh.GetNumberOfCells()):
            cell = mesh.GetCell(i)
            cell_points = [cell.GetPointId(j) for j in range(cell.GetNumberOfPoints())]
            cells.append(cell_points)
        
        return {
            'points': points,
            'cells': cells,
            'vtk_mesh': mesh  # Keep the VTK object for later use
        }
    
    @staticmethod
    def _read_stl(file_path):
        """Read STL file using VTK"""
        if not VTK_AVAILABLE:
            raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        
        reader = vtk.vtkSTLReader()
        reader.SetFileName(file_path)
        reader.Update()
        
        mesh = reader.GetOutput()
        
        # Extract points
        points = vtk_to_numpy(mesh.GetPoints().GetData())
        
        # Extract cells (triangles in STL)
        cells = []
        for i in range(mesh.GetNumberOfCells()):
            cell = mesh.GetCell(i)
            cell_points = [cell.GetPointId(j) for j in range(cell.GetNumberOfPoints())]
            cells.append(cell_points)
        
        return {
            'points': points,
            'cells': cells,
            'vtk_mesh': mesh  # Keep the VTK object for later use
        }
    
    @staticmethod
    def _read_stl_binary(file_path):
        """Read binary STL file without VTK"""
        with open(file_path, 'rb') as f:
            # Skip header
            f.seek(80)
            
            # Read number of triangles
            num_triangles = struct.unpack('<I', f.read(4))[0]
            
            # Initialize arrays
            points = []
            cells = []
            
            # Read triangles
            for _ in range(num_triangles):
                # Skip normal vector (3 floats)
                f.seek(12, 1)
                
                # Read vertices
                v1 = struct.unpack('<3f', f.read(12))
                v2 = struct.unpack('<3f', f.read(12))
                v3 = struct.unpack('<3f', f.read(12))
                
                # Skip attribute byte count
                f.seek(2, 1)
                
                # Add vertices to points array
                points.extend([v1, v2, v3])
                
                # Add cell (triangle)
                base_idx = len(points) - 3
                cells.append([base_idx, base_idx + 1, base_idx + 2])
        
        # Convert to numpy arrays
        points = np.array(points)
        cells = np.array(cells)
        
        return {
            'points': points,
            'cells': cells
        }
    
    @staticmethod
    def _write_vtk(mesh_data, file_path):
        """Write mesh to VTK file"""
        if not VTK_AVAILABLE:
            raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        
        if 'vtk_mesh' in mesh_data:
            # If we have the original VTK mesh, use it
            mesh = mesh_data['vtk_mesh']
        else:
            # Create a new VTK mesh
            points = vtk.vtkPoints()
            for point in mesh_data['points']:
                points.InsertNextPoint(point)
            
            cells = vtk.vtkCellArray()
            for cell in mesh_data['cells']:
                vtk_cell = vtk.vtkIdList()
                for point_id in cell:
                    vtk_cell.InsertNextId(point_id)
                cells.InsertNextCell(vtk_cell)
            
            mesh = vtk.vtkUnstructuredGrid()
            mesh.SetPoints(points)
            mesh.SetCells(vtk.VTK_TRIANGLE, cells)
        
        writer = vtk.vtkUnstructuredGridWriter()
        writer.SetFileName(file_path)
        writer.SetInputData(mesh)
        writer.Write()
    
    @staticmethod
    def _write_stl(mesh_data, file_path):
        """Write mesh to STL file using VTK"""
        if not VTK_AVAILABLE:
            raise ImportError("VTK is not available. Install it with 'pip install vtk'")
        
        if 'vtk_mesh' in mesh_data:
            # If we have the original VTK mesh, use it
            mesh = mesh_data['vtk_mesh']
        else:
            # Create a new VTK mesh
            points = vtk.vtkPoints()
            for point in mesh_data['points']:
                points.InsertNextPoint(point)
            
            cells = vtk.vtkCellArray()
            for cell in mesh_data['cells']:
                if len(cell) == 3:  # STL only supports triangles
                    triangle = vtk.vtkTriangle()
                    triangle.GetPointIds().SetId(0, cell[0])
                    triangle.GetPointIds().SetId(1, cell[1])
                    triangle.GetPointIds().SetId(2, cell[2])
                    cells.InsertNextCell(triangle)
            
            mesh = vtk.vtkPolyData()
            mesh.SetPoints(points)
            mesh.SetPolys(cells)
        
        writer = vtk.vtkSTLWriter()
        writer.SetFileName(file_path)
        writer.SetInputData(mesh)
        writer.Write()
    
    @staticmethod
    def _write_stl_binary(mesh_data, file_path):
        """Write mesh to binary STL file without VTK"""
        points = mesh_data['points']
        cells = mesh_data['cells']
        
        with open(file_path, 'wb') as f:
            # Write header
            f.write(b'\x00' * 80)
            
            # Write number of triangles
            f.write(struct.pack('<I', len(cells)))
            
            # Write triangles
            for cell in cells:
                if len(cell) != 3:
                    continue  # Skip non-triangular cells
                
                # Get vertices
                v1 = points[cell[0]]
                v2 = points[cell[1]]
                v3 = points[cell[2]]
                
                # Calculate normal
                normal = np.cross(v2 - v1, v3 - v1)
                norm = np.linalg.norm(normal)
                if norm > 0:
                    normal = normal / norm
                else:
                    normal = np.array([0, 0, 1])
                
                # Write normal
                f.write(struct.pack('<3f', *normal))
                
                # Write vertices
                f.write(struct.pack('<3f', *v1))
                f.write(struct.pack('<3f', *v2))
                f.write(struct.pack('<3f', *v3))
                
                # Write attribute byte count
                f.write(struct.pack('<H', 0)) 